# -*- coding: utf-8 -*-
# @Time    : 2024/9/12 13:32
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : gMOT_4beam.py
# @Software: PyCharm

import matplotlib.pyplot as plt
import numpy
from mpl_toolkits.mplot3d.axes3d import Axes3D

from mot.atom import Rb87_D2
from mot.average_scattering_force_model import MOT
from mot.field.base import RectGrid3D
from mot.field.laser import NonDiffractionLaserBeam, LaserInfo, BeamPolarization, set_frequency
from mot.field.magnetic import BGenerators
from mot.field.magnetic import CurrentCarryingLoop
from mot.gmot.grating import Sector, GratingSection, LaserDiffractedByGratingSection
from myutils.geometrical import Plane
from _logging import logger

def f(aperture,intensity_of_incident_laser,phi,reflectivity,z0,I_coil,coil_R,coil_theta=0,coil_phi=0,coil_x=0,coil_y=0,coil_z=-100e-3):
    """
    aperture ： float
        光束半径 m
    intensity_of_incident_laser： float
        入射激光强度 W/m^2
    phi : numpy.float64
        光栅衍射角度
    reflectivity :float
        反射率
    z0 : float
        四极中心位置 m
    I_coil : float
        线圈电流 A
    coil_R : float
        线圈半径 m
    coil_theta : numpy.float64
        线圈方位角
    coil_phi : numpy.float64
        线圈俯仰角
    coil_x : float
        线圈中心x坐标 m
    coil_y : float
        线圈中心y坐标 m
    coil_z : float
        线圈中心z坐标 m
    """
    # f(30e-3 / 2, (17.3 * 0 + 16.8) * 1e-3 / (1e-2) ** 2, numpy.deg2rad(42.5 * 0 + 46.3),36.3e-2,4.6e-3,229.71598951823216 * 16.1,150e-3,0,0, 0,0, -100e-3)
    grid = RectGrid3D(numpy.array((*numpy.linspace(-10e-3, -5.5e-3, 10),
                                   *numpy.linspace(-5e-3, 5e-3, 25),
                                   *numpy.linspace(5.5e-3, 10e-3, 10)
                                   )),
                      numpy.array((*numpy.linspace(-10e-3, -5.5e-3, 10),
                                   *numpy.linspace(-5e-3, 5e-3, 26),
                                   *numpy.linspace(5.5e-3, 10e-3, 10)
                                   )),
                      numpy.array((*numpy.linspace(0, 8e-3, 25),
                                   *numpy.linspace(9e-3, 15e-3, 15)
                                   )
                                  ),
                      [
                          # "Bx", "By", "Bz"
                      ]
                      )
    # grid = RectGrid3D(numpy.array((*numpy.linspace(-10e-3, 12e-3, 35),
    #                                )),
    #                   numpy.array((*numpy.linspace(-10e-3, 12e-3, 36),
    #                                )),
    #                   numpy.array((*numpy.linspace(0, 13e-3, 40),
    #                                )
    #                               ),
    #                   [
    #                       # "Bx", "By", "Bz"
    #                   ]
    #                   )
    grating_sections = [GratingSection(
        Sector(numpy.array([0, 0]), 10e-3, 1e-3, numpy.deg2rad(-60 + i * 120), numpy.deg2rad(60 + i * 120)).shape,
        Plane(numpy.array([0, 0, grid.zs.min()]), numpy.array([0, 0, 1]))) for i in range(3)
    ]
    linedata = Rb87_D2()
    detuning_Hz = -1.61 * linedata.Gamma  # -1.72 * linedata.Gamma
    __polarizations = [+1, -1]
    incident_laser = NonDiffractionLaserBeam(
        set_frequency(LaserInfo(numpy.array([0, 0, -1]), BeamPolarization(__polarizations[0]), "incident", detuning_Hz),
                      linedata.f_0 + detuning_Hz),
        Plane(numpy.array([0, 0, grid.zs.max()]), numpy.array([0, 0, 1]), ),
        lambda x, y, z: intensity_of_incident_laser if x ** 2 + y ** 2 < aperture ** 2 else 0.,
    )
    lasers_diffracted_by_grating_section = [
        LaserDiffractedByGratingSection(
            grating_sections[i], incident_laser,
            set_frequency(LaserInfo(numpy.array(  # [-1, 0, 2]
                [*(1 * numpy.array([numpy.cos(numpy.deg2rad(180 + 120 * i)),
                                    numpy.sin(numpy.deg2rad(180 + 120 * i))])),
                 numpy.tan(phi)]), BeamPolarization(__polarizations[1]),
                "grating%d" % i, detuning_Hz), linedata.f_0 + detuning_Hz), reflectivity / numpy.sin(phi)
        ) for i, gs in
        enumerate(grating_sections)]
    coil1_direction=numpy.array([numpy.cos(coil_theta)*numpy.sin(coil_phi), numpy.sin(coil_theta)*numpy.sin(coil_phi),numpy.cos(coil_phi)])
    coil2_direction=[i * -1 for i in coil1_direction]
    """
    coil_direction 线圈法向量
    """
    coil1_center_position=numpy.array([coil_x, coil_y, coil_z])
    coil2_center_position=numpy.array([-coil_x, -coil_y, 2 * z0 - coil_z])
    coil1 = CurrentCarryingLoop(I_coil, coil_R, coil1_center_position,coil1_direction)
    coil2 = CurrentCarryingLoop(I_coil, coil_R, coil2_center_position, coil2_direction)
    mot_initializer = MOT(grid, BGenerators([coil1, coil2]), [incident_laser, *lasers_diffracted_by_grating_section],
                          linedata)
    mot_initializer.cache()

    ts_dumped = numpy.linspace(0, 30e-3, 200)
    from scipy.integrate import solve_ivp


    def dY_dt(t, Y_: numpy.ndarray):
        """

        :param t:
        :param Y: shape (6*N_particles, ), 其中每6位分别为(x, y, z, vx, vy, vz)
        :return:
        """
        Y = Y_.reshape((-1, 6))
        ret = numpy.zeros(Y.shape)
        ret[:, :3] = Y[:, -3:]
        ret[:, -3:] = mot_initializer.scattering_force_from_laser(
            *[Y[..., i] for i in range(6)]) / mot_initializer.transition_line_data.mass_kg
        return ret.reshape((-1,))


    # trajdata = solve_ivp(dY_dt, (ts_dumped[0], ts_dumped[-1]), numpy.array([
    #     [0, 1e-3, 1e-3, 1, 1, 1],
    #     [1e-3, 1e-3, 5e-3, 1, -1, 1],
    # ]).T, t_eval=ts_dumped,vectorized  = True)
    N_0 = int(1e4)
    position_0 = numpy.array([grid.xs.min(), grid.ys.min(), grid.zs.min()]) + (
            numpy.array([grid.xs.max(), grid.ys.max(), grid.zs.max()]) -
            numpy.array([grid.xs.min(), grid.ys.min(), grid.zs.min()])) * numpy.random.random(
        (N_0, 3))  # grid.all_points().reshape((-1, 3))
    i_particles = numpy.arange(len(position_0))[numpy.random.random(len(position_0)) > 0.9]
    position_0 = position_0[i_particles]
    v_0 = numpy.zeros(position_0.shape)
    Y0 = numpy.array((*position_0.transpose((1, 0)), *v_0.transpose((1, 0)))).transpose((1, 0))
    from mot.particle_initializer import MaxwellParticleInitializer

    min_and_max = lambda xs: numpy.array((xs.min(), xs.max()))

    Y0 = MaxwellParticleInitializer(
        N_0,
        numpy.array([min_and_max(grid.xs),
                     min_and_max(grid.ys),
                     min_and_max(grid.zs),
                     ]), 0, 20, 1.0, mot_initializer.transition_line_data.mass_kg).generate().values

    trajdata = solve_ivp(dY_dt, (ts_dumped[0], ts_dumped[-1]), y0=Y0.reshape((-1,)), t_eval=ts_dumped, vectorized=False)
    N_particles = int(len(trajdata.y[:, 0]) / 6)

    MOT_center = numpy.array([0, 0, 8e-3])
    last_pos = trajdata.y.reshape((N_particles, 6, -1))[:, :3, -1]
    last_vel = trajdata.y.reshape((N_particles, 6, -1))[:, -3:, -1]
    indexes_of_trapped_particles = (numpy.linalg.norm(last_pos - MOT_center, axis=1) < 3e-3) & (
            numpy.linalg.norm(last_vel, axis=1) < 10)

    capture_rate = len(last_pos[indexes_of_trapped_particles]) / N_particles
    # logger.info(" 光束半径：%.8f" % (aperture))
    # logger.info(" 入射激光强度：%.8f" % (intensity_of_incident_laser))
    # logger.info(" 光栅衍射角度：%.8f" % (phi))
    # logger.info(" 反射率：%.8f" % (reflectivity))
    # logger.info(" 四极中心位置：%.8f" % (z0))
    # logger.info(" 线圈电流：%.8f" % (I_coil))
    # logger.info(" 线圈半径：%.8f" % (coil_R))
    # logger.info(" 线圈方位角：%.8f" % (coil_theta))
    # logger.info(" 线圈俯仰角：%.8f" % (coil_phi))
    # logger.info(" 线圈中心x坐标：%.8f" % (coil_x))
    # logger.info(" 线圈中心y坐标：%.8f" % (coil_y))
    # logger.info(" 线圈中心z坐标：%.8f" % (coil_z))
    logger.info(" 俘获效率：%.2f%%" % (100 * capture_rate))
    return capture_rate